package classes;
import java.util.*;
import java.io.*;


public class IntexSentence {

  // sgml IDs
  static public String STR_P1 = "p1";
  static public String STR_P2 = "p2";
  static public String STR_PAIR = "pair";

  // ACE sentence
  ACETokSentence m_aceTS;
  // array of sentence entities
  ACEEMention[] m_ems;
  // upper-triangular matrix of entity pairs
  EntityPair[][] m_pairs;


  public IntexSentence(ACETokSentence aceTS)
  {
    m_aceTS = aceTS;

    m_ems = (ACEEMention[]) aceTS.m_lstEM.toArray(new ACEEMention[0]);
    // sort ementions based on head position
    Arrays.sort(m_ems);

    // create matrix of entity pairs
    m_pairs = new EntityPair[m_ems.length][m_ems.length];
    
    // populate matrix with positive pairs
    ListIterator it = aceTS.m_lstRM.listIterator();
    while (it.hasNext()) {
      ACERMention rm = (ACERMention) it.next();
      int nPos1 = findEntity(rm.m_arg1);
      int nPos2 = findEntity(rm.m_arg2);
      EntityPair pair;
      int nDir;
      if (nPos1 <= nPos2) {
	m_pairs[nPos1][nPos2] = pair = new EntityPair(rm.m_strID, nPos1, nPos2);
	nDir = 0;
      }
      else {
	m_pairs[nPos2][nPos1] = pair = new EntityPair(rm.m_strID, nPos2, nPos1);
	nDir = 1;
      }
      String strType = rm.m_relation.m_strType;
      String strSType = rm.m_relation.m_strSType;
      if (!EntityPair.Labels.subtype && strType.equals(EntityPair.Labels.RELATION)) {
    	  pair.m_nTrueLabel = EntityPair.Labels.RELATION_SYM;
      }
      else if (EntityPair.Labels.subtype && strSType.equals(EntityPair.Labels.RELATION)){
    	  
    	  pair.m_nTrueLabel = EntityPair.Labels.RELATION_SYM;
    	  
      } else
    	  
    	  assert false;
    }
    
    // complete matrix with negative entity pairs
    for (int i = 0; i < m_ems.length; i++) {
      for (int j = i + 1; j < m_ems.length; j++) {
	if (m_pairs[i][j] == null) {
	  m_pairs[i][j] = new EntityPair("", i, j);
	  m_pairs[i][j].m_nTrueLabel = EntityPair.Labels.NONE;
	}
	if (m_ems[i].includes(m_ems[j]) || m_ems[j].includes(m_ems[i]))
	  m_pairs[i][j] = null;
      }
    }
  }


  protected int findEntity(ACEEMention em)
  {
    for (int i = 0; i < m_ems.length; i++)
      if (m_ems[i] == em)
	return i;
    
    assert false;
    return -1;
  }
  

  public int[] pairsCounts()
  {
    int[] counts = new int[2];
    for (int i = 0; i < m_ems.length; i++)
      for (int j = i + 1; j < m_ems.length; j++)
	if (m_pairs[i][j] != null)
	  if (m_pairs[i][j].m_nTrueLabel == EntityPair.Labels.NONE)
	    counts[0]++;
	  else
	    counts[1]++;
    
    return counts;
  }


  public int[] getStats()
  {
    int[] stats = new int[3];
    for (int i = 0; i < m_ems.length; i++)
      for (int j = i + 1; j < m_ems.length; j++)
	if (m_pairs[i][j] != null) {
	  if (m_pairs[i][j].m_nTrueLabel != EntityPair.Labels.NONE)
	    stats[0]++;
	  if (m_pairs[i][j].m_nInfLabel != EntityPair.Labels.NONE)
	    stats[1]++;
	  if (m_pairs[i][j].m_nInfLabel == m_pairs[i][j].m_nTrueLabel &&
	      m_pairs[i][j].m_nInfLabel != EntityPair.Labels.NONE)
	    stats[2]++;
	}
    
    return stats;
  }


  public int writeInf(PrintWriter pw, int crtPairs)
  {
    int nPairs = 0;
   
    LinkedList lstAnno = new LinkedList();
    // add entities
    for (int i = 0; i < m_ems.length; i++) {
      ACEEMention anno = new ACEEMention();
      anno.m_nBeg = m_ems[i].m_nBeg;
      anno.m_nLen = m_ems[i].m_nLen;
      anno.m_strID = "prot";
      lstAnno.add(anno);
    }

    // add inferred pairs
    for (int i = 0; i < m_ems.length; i++)
      for (int j = i + 1; j < m_ems.length; j++) {
        if (m_pairs[i][j] != null/* && m_pairs[i][j].isInfPos()*/) {
          double conf = 0.0;
          if (m_pairs[i][j].m_nInfLabel == 1)
            conf = m_pairs[i][j].m_conf;
          else
            conf = 1 - m_pairs[i][j].m_conf;
          if (conf > 0.1) {
            // add annotation for first element of the pair
            ACEEMention anno1 = new ACEEMention();
	    anno1.m_nBeg = m_ems[i].m_nBeg;
	    anno1.m_nLen = m_ems[i].m_nLen;
	    anno1.m_strID = STR_P1;

            anno1.setAttribute(STR_PAIR, Integer.toString(crtPairs + nPairs));
            anno1.setAttribute("conf", Double.toString(conf));
            lstAnno.add(anno1);
            // add annotation for second element of the pair
            ACEEMention anno2 = new ACEEMention();
	    anno2.m_nBeg = m_ems[j].m_nBeg;
	    anno2.m_nLen = m_ems[j].m_nLen;
	    anno2.m_strID = STR_P2;

            anno2.setAttribute(STR_PAIR, Integer.toString(crtPairs + nPairs));
            anno2.setAttribute("conf", Double.toString(conf));
            lstAnno.add(anno2);
            // increment number of pairs
            nPairs++;
          }
        }
      }
    
    // write annotations to stream
    m_aceTS.writeAnnotations(lstAnno, pw);
    
    return nPairs;
  }

}
